import tvm
from tvm.contrib import graph_executor
import numpy as np


def run():
    dev = tvm.cpu(0)
    lib = tvm.runtime.load_module("./data/model_results/VGG19.h5.so")
    original_model = graph_executor.GraphModule(lib["default"](dev))
    # img = np.random.randn(1, 224, 224, 3)
    shape = (1, 224, 224, 3)
    img = np.ones(shape)
    dtype = "float32"
    img = img.astype(dtype)
    original_model.set_input("input_3", img)
    original_model.run()

    output_count = original_model.get_num_outputs()
    out = [original_model.get_output(i).asnumpy() for i in range(output_count)]
    print(out[0])


if __name__ == "__main__":
    run()
